from sqlalchemy import Column, String, DateTime
from sqlalchemy.orm import sessionmaker

from init.mysql import Base, engine


class StockInfo(Base):
    __tablename__ = 'stock_info'
    dm = Column(String(6), primary_key=True)
    mc = Column(String(255))
    jys = Column(String(2))

    def __init__(self, dm, mc, jys):
        self.dm = dm
        self.mc = mc
        self.jys = jys
        self.update_time = Column(DateTime)


Base.metadata.create_all(engine)


def insert(session, stock_info):
    session.add(stock_info)
    session.commit()


def update(session, stock_info):
    session.query(StockInfo).filter(StockInfo.id == stock_info.id).update({
        StockInfo.dm: stock_info.dm,
        StockInfo.mc: stock_info.mc,
        StockInfo.jys: stock_info.jys
    })
    session.commit()


def find_all(session):
    return session.query(StockInfo).all()


def batch_insert_or_update(data_list):
    batch_size = 1000
    i = 0
    session = sessionmaker(bind=engine)()
    try:
        for data in data_list:
            # 查询是否已存在相同的记录
            existing_record = session.query(StockInfo).filter_by(dm=data['dm']).first()
            if existing_record:
                # 如果存在，更新记录
                existing_record.mc = data['mc']
                existing_record.jys = data['jys']
            else:
                # 如果不存在，插入新记录
                new_record = StockInfo(dm=data['dm'], mc=data['mc'], jys=data['jys'])
                session.add(new_record)
            i += 1
            if i % batch_size == 0:
                i = 0
                session.commit()
        session.commit()
    except Exception as e:
        print(f"Error: {e}")
        session.rollback()
    finally:
        session.close()
